import os, sys

parent_path = os.path.abspath(os.path.join(__file__, *(['..'] * 5)))
sys.path.insert(0, parent_path)

# from queue import Queue
from core.task.modules.processors.base_processor import BaseProcessor
from threading import Thread
import multiprocessing as mp
from typing import *
import cv2, time, traceback
import numpy as np
from core.utils.logger import logging
from core.utils.error_exception import ErrorException


class SteamPull(BaseProcessor):
    """任务管道的拉流组件
    跟随任务启动后, 会在task_init时指定拉流地址"""
    def __init__(self, keys, 
                 SteamPull_url = None, 
                 SteamPull_size = None, 
                 SteamPull_fps = 200,
                 SteamPull_retry_time = 5) -> None:
        """
        input:
            keys:{
                out: 在流转数据box中的键名
            }
            pull_url: 拉流地址
            pull_size: 拉流图像的大小
            retry_time: 拉流时重试的次数
        """
        super().__init__(keys)
        self.keys = keys
        self.pull_size = SteamPull_size
        self.pull_url = SteamPull_url
        self.cap = None
        self.frame = None
        self.fps = SteamPull_fps
        self.retry_time = SteamPull_retry_time
        self.run = None
    
    def _init_check(self):
        assert not self.pull_url is None
        assert not self.cap is None
        assert not self.pull_size is None
        assert not self.run is None

    def _set(self, data:Dict):
        """
        当本类作为module被Task Pipeliner调用时起效, 作为任务初始化函数被调用
        """
        if data.get("SteamPull_url",False):
            self.pull_url = data["SteamPull_url"]
        if data.get("SteamPull_size",False):
            self.pull_size = data["SteamPull_size"]
        self.cap, self.frame = self.init_cap(self.pull_url)
        self.run = Thread(target=self.video_read_loop)
    
    def init_cap(self, pull_url):
        cap = cv2.VideoCapture(pull_url)
        # assert cap.isOpened(), "{} can't open".format(pull_url)
        if not cap.isOpened():
            raise ErrorException(f"Failed to open stream: {pull_url}")
        cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
        ret, frame = cap.read()
        assert ret, "{} can't open".format(pull_url)
        return cap, frame

    def video_read_loop(self):
        try:
            while True:
                time.sleep(1)
                if self.is_running():
                    break
                elif self.is_close():
                    break
                logging.info("{} waitting start....".format(self.__class__.__name__))

            retry_time = 0
            while self.is_running():
                ret, frame = self.cap.read()
                # 用于处理空帧
                if not ret:
                    ret, frame = self.cap.read()
                    if not ret:
                        ret, frame = self.cap.read()
                # ------
                if not ret:
                    retry_time+=1
                    if retry_time > self.retry_time:
                        raise ErrorException("cap retry more than {}, close this cap".format(self.retry_time))
                    logging.warning("|WAR: cap reading fail, waiting 2s and retry {}|".format(retry_time))
                    time.sleep(2)
                    self.cap = cv2.VideoCapture(self.pull_url)
                    continue
                retry_time = 0
                self.frame = cv2.resize(frame,self.pull_size)
                time.sleep(1/self.fps)

        except Exception:
            logging.error(traceback.format_exc())
            self.close()
        finally:
            logging.info("|SteamPull.video_read_loop done|")
    
    def _start(self):
        # 启动循环
        self.run.start()

    def _call(self, data:Dict) -> Dict:
        data[self.keys["out"]]= self.frame[None,...]
        return data

class SteamPullMP(BaseProcessor):
    """任务管道的拉流组件
    跟随任务启动后, 会在 task_init 时指定拉流地址"""
    
    def __init__(self, keys, 
                 SteamPull_url=None, 
                 SteamPull_size=None, 
                 SteamPull_fps=200,
                 SteamPull_retry_time=5) -> None:
        """
        input:
            keys:{
                out: 在流转数据 box 中的键名
            }
            pull_url: 拉流地址
            pull_size: 拉流图像的大小 (宽, 高)
            retry_time: 拉流时重试的次数
        """
        super().__init__(keys)
        self.keys = keys
        self.pull_size = SteamPull_size
        self.pull_url = SteamPull_url
        self.frame = None
        self.fps = SteamPull_fps
        self.retry_time = SteamPull_retry_time
        self.process = None
        self.frame_queue = mp.Queue(maxsize=1)  # 用于进程间通信的队列
        self.control_queue = mp.Queue(maxsize=1)  # 用于控制进程状态

    def _init_check(self):
        assert self.pull_url is not None
        assert self.pull_size is not None

    def _set(self, data: Dict):
        """
        当本类作为 module 被 Task Pipeliner 调用时起效, 作为任务初始化函数被调用
        """
        if data.get("SteamPull_url", False):
            self.pull_url = data["SteamPull_url"]
        if data.get("SteamPull_size", False):
            self.pull_size = data["SteamPull_size"]
        _, self.frame = self.init_cap(self.pull_url)

        # 启动视频读取进程
        self.process = mp.Process(target=self.video_read_loop, args=(self.pull_url, self.pull_size, self.frame_queue, self.control_queue))
        self.process.daemon = True

    def init_cap(self, pull_url):
        cap = cv2.VideoCapture(pull_url)
        assert cap.isOpened(), "{} can't open".format(pull_url)
        cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
        ret, frame = cap.read()
        assert ret, "{} can't open".format(pull_url)
        return cap, frame

    def video_read_loop(self, pull_url, pull_size, frame_queue, control_queue):
        """
        独立进程中运行的视频读取循环
        """
        try:
            cap = cv2.VideoCapture(pull_url)
            if not cap.isOpened():
                raise ErrorException(f"Failed to open stream: {pull_url}")
            cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)

            retry_time = 0
            while True:
                # 检查控制队列是否有退出信号
                if not control_queue.empty():
                    command = control_queue.get()
                    if command == "STOP":
                        logging.info("Video reading process received STOP signal.")
                        break

                # 读取视频帧
                ret, frame = cap.read()
                if not ret:
                    retry_time += 1
                    if retry_time > self.retry_time:
                        raise ErrorException(f"Failed to read frame after {self.retry_time} retries, closing stream.")
                    logging.warning(f"Retrying to read frame ({retry_time}/{self.retry_time})...")
                    time.sleep(2)
                    cap = cv2.VideoCapture(pull_url)
                    continue

                # 重置重试次数
                retry_time = 0

                # 调整帧大小
                frame = cv2.resize(frame, pull_size)
                try:
                    frame_queue.put_nowait(frame)
                except:
                    pass
                # # 将帧放入队列，非阻塞地覆盖旧帧
                # if not frame_queue.full():
                #     frame_queue.put(frame)
                # else:
                #     frame_queue.get()  # 丢弃旧帧
                #     frame_queue.put(frame)

                # 控制帧率
                time.sleep(1 / self.fps)

        except Exception as e:
            logging.error(f"Exception in video_read_loop: {traceback.format_exc()}")
        finally:
            cap.release()
            logging.info("Video reading process terminated.")

    def _start(self):
        # 启动视频读取进程
        if self.process:
            self.process.start()

    def _call(self, data: Dict) -> Dict:
        """
        从队列中获取最新帧，并返回给数据流
        """
        if not self.frame_queue.empty():
            self.frame = self.frame_queue.get()  # 获取最新帧
        data[self.keys["out"]] = self.frame[None, ...]
        return data

    def _close(self):
        """
        关闭视频读取进程
        """
        if self.process and self.process.is_alive():
            self.control_queue.put("STOP")  # 发送停止信号
            self.process.join()  # 等待进程结束

class SteamPulls(BaseProcessor):
    """任务管道的拉流组件
    跟随任务启动后, 会在task_init时指定拉流地址"""
    def __init__(self, keys, 
                 SteamPull_urls = None, 
                 SteamPull_size = None, 
                 SteamPull_retry_time = 5) -> None:
        """
        input:
            keys:{
                out: 在流转数据box中的键名
            }
            pull_url: 拉流地址
            pull_size: 拉流图像的大小
            retry_time: 拉流时重试的次数
        """
        super().__init__(keys)
        self.keys = keys
        self.pull_size = SteamPull_size
        self.pull_urls = SteamPull_urls
        self.caps = []
        self.frames = None
        self.retry_time = SteamPull_retry_time
        self.run = None
    
    def _init_check(self):
        assert not self.pull_urls is None
        assert len(self.caps) > 0
        assert not self.pull_size is None
        assert not self.run is None

    def _set(self, data:Dict):
        """
        当本类作为module被Task Pipeliner调用时起效, 作为任务初始化函数被调用
        """
        if data.get("SteamPull_url",False):
            self.pull_urls = data["SteamPull_url"]
        if data.get("SteamPull_size",False):
            self.pull_size = data["SteamPull_size"]
        self.caps, self.frames = self.init_cap(self.pull_urls)
        self.run = Thread(target=self.video_read_loop)
    
    def _init_cap(self, pull_url):
        cap = cv2.VideoCapture(pull_url)
        assert cap.isOpened(), "{} can't open".format(pull_url)
        ret, frame = cap.read()
        assert ret, "{} can't open".format(pull_url)
        return cap, frame
    
    def init_cap(self, pull_urls):
        frames = []
        caps = []
        for pull_url in pull_urls:
            cap, frame = self._init_cap(pull_url)
            caps.append(cap)
            frame = cv2.resize(frame,self.pull_size)
            frames.append(frame)
        return caps, np.array(frames,dtype=np.uint8)

    def video_read_loop(self):
        try:
            while True:
                time.sleep(1)
                if self.is_running():
                    break
                elif self.is_close():
                    break
                logging.info("{} waitting start....".format(self.__class__.__name__))

            while self.is_running():
                for idx, cap in enumerate(self.caps):
                    ret, frame = cap.read()
                    if not ret:
                        pass
                    frame = cv2.resize(frame,self.pull_size)
                    self.frames[idx] = frame
                
                time.sleep(1/200)

        except Exception:
            logging.error(traceback.format_exc())
            self.close()
        finally:
            logging.info("|SteamPulls.video_read_loop done|")
    
    def _start(self):
        # 启动循环
        self.run.start()

    def _call(self, data:Dict) -> Dict:
        data[self.keys["out"]]= self.frames
        return data


if __name__ == "__main__":
    def multi_steam_test():
        steam_puller = SteamPulls(
            keys={
                "out":"pull_data"
                },
            SteamPull_urls = ["rtmp://218.94.155.86:1936/live/livestream_74_1",
                            "rtmp://218.94.155.86:1936/live/livestream_74_2",],
            SteamPull_size=[1280,720],
        )
        data = {
        }
        steam_puller.set(data)
        steam_puller.start()
        idx = 0
        while idx<20:
            data = steam_puller(data)
            cv2.imwrite("temp1.jpg",data["pull_data"][0])
            cv2.imwrite("temp2.jpg",data["pull_data"][1])
            time.sleep(1)
            # idx+=1
        steam_puller.close()

    def single_steam_test():
        steam_puller = SteamPullMP(
            keys={
                "out":"pull_data"
                },
            SteamPull_url = "rtmp://218.94.155.86:1936/live/livestream_74_1",
            SteamPull_size=[1280,720],
        )
        data = {
        }
        steam_puller.set(data)
        steam_puller.start()
        idx = 0
        while idx<20:
            data = steam_puller(data)
            cv2.imwrite("temp1.jpg",data["pull_data"][0])
            time.sleep(1)
            # idx+=1
        steam_puller.close()
        print("___the other task___")
    single_steam_test()
    # multi_steam_test()
    
